import json
import os
import sys
from PIL import Image
import cv2
import numpy as np


def update(num):
    """
    空函数：裁切并保存完一张图像就调用此函数
    :return:
    """
    sys.stdout.write(str(num))
    sys.stdout.flush()


def generate_trapezoidal_image(input_path, tilt_angle, output_path):
    """
    矩形变换为梯形
    :param input_path: 所变换图像的路径
    :param tilt_angle: 变换倾角
    :param output_path: 变换后图像的保存路径
    """
    image = cv2.imread(input_path)
    width, height = image.shape[1], image.shape[0]

    # 根据倾角计算目标梯形的四个角点坐标
    half_tan_angle = np.tan(np.radians(tilt_angle) / 2)
    top_width = width - 2 * half_tan_angle * height
    dst_points = np.float32([
        [half_tan_angle * height, 0],
        [width - half_tan_angle * height, 0],
        [0, height],
        [width, height]
    ])

    # 计算透视变换矩阵
    matrix = cv2.getPerspectiveTransform(np.float32([[0, 0], [width - 1, 0], [0, height - 1], [width - 1, height - 1]]),
                                         dst_points)

    # 进行透视变换
    warped_image = cv2.warpPerspective(
        image, matrix, (width, height), flags=cv2.INTER_LINEAR)

    # 保存变换后的图像
    cv2.imwrite(output_path, warped_image)


def crop_images(coord_list, width=200, height=200, original_image_path='', output_folder=''):
    """
    裁切图片并保存到指定文件夹中
    :param coord_list: 左上角裁切坐标的列表
    :param width:
    :param height:
    :param original_image_path: 原图的路径
    :param output_folder: 裁切后图像保存的路径
    :return:
    """
    image = Image.open(original_image_path)

    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    for index, coords in enumerate(coord_list):
        left, top = coords
        right, bottom = left + width, top + height

        cropped_image = image.crop((left, top, right, bottom))

        output_image_name = f'{index + 1}.jpg'
        output_path = os.path.join(output_folder, output_image_name)

        # 保存裁切后的图像
        cropped_image.save(output_path)
        update(index + 1)


def scan(trajectory_path, angle, save_crop_imgs_dir, img_path) -> None:
    """

    :param trajectory_path: 裁切坐标JSON文件
    :param angle: 卫星拍摄倾角
    :param save_crop_imgs_dir: 裁切后的保存文件夹路径
    :param img_path: 原图路径
    :return:
    """

    # 1.获取裁切坐标
    with open(trajectory_path, 'r') as file:
        data = json.load(file)

    # 获取路径中JSON文件的名字，便于读取JSON文件中的数据
    value_name = os.path.basename(trajectory_path).split(
        '.')[-2]  # trajectory_diagonal
    trajectory_data_lst = data.get(value_name)

    # 2.根据坐标裁切图片并保存到指定文件夹下
    crop_images(trajectory_data_lst, width=300, height=300,
                original_image_path=img_path,
                output_folder=save_crop_imgs_dir)

    # 3. 根据卫星的倾角调整图像
    jpg_files = [os.path.abspath(os.path.join(save_crop_imgs_dir, filename)) for filename in os.listdir(save_crop_imgs_dir) if
                 filename.endswith('.jpg')]

    for jpg_file in jpg_files:
        generate_trapezoidal_image(jpg_file, angle / 2, jpg_file)


if __name__ == '__main__':
    """
    将原大图裁切成 200 * 200 的小图
    """
    import argparse

    parser = argparse.ArgumentParser(description="执行图像裁切操作")
    parser.add_argument("--traj", type=str)
    parser.add_argument("--angle", type=int)
    args = parser.parse_args()

    # 轨迹路径的JSON文件

    trajectory_path = r'./src/trajs/trajectory_%s.json' % args.traj
    # 裁切后图像保存的文件夹路径
    save_crop_imgs_dir = r'./output_images'
    # 原图路径
    img_path = r'./src/global-view.jpg'  # 原图路径

    scan(trajectory_path, args.angle, save_crop_imgs_dir, img_path)
